"""
Event type definitions for the continual learning framework.

This module defines the event types that components can listen for and emit.
Each event type carries specific data relevant to that event.
"""

from enum import Enum, auto
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, TypeVar, Generic, Union
import time
from torch.utils.data import DataLoader
import pytorch_lightning as pl


class EventType(Enum):
    """Enumeration of event types in the continual learning framework."""

    # Task lifecycle events
    TASK_START = auto()
    TASK_END = auto()

    # Data events
    BATCH_AVAILABLE = auto()
    DATASET_AVAILABLE = auto()

    # Training lifecycle events
    EPOCH_START = auto()
    EPOCH_END = auto()
    BATCH_START = auto()
    BATCH_END = auto()
    TRAINING_START = auto()
    TRAINING_END = auto()

    # Evaluation events
    EVALUATION_START = auto()
    EVALUATION_END = auto()
    METRICS_COMPUTED = auto()

    # Memory events
    MEMORY_UPDATE = auto()

    # Model events
    MODEL_SAVED = auto()
    MODEL_LOADED = auto()
    PARAMETER_UPDATE = auto()

    # Logging events
    LOG_METRICS = auto()


@dataclass
class EventData:
    """Base class for all event data."""

    event_type: EventType
    timestamp: float = field(default_factory=lambda: time.time())
    source: Optional[str] = None


@dataclass
class TaskEvent(EventData):
    """Event data for task-related events."""

    task_id: int = 0
    pl_module: Optional["pl.LightningModule"] = None
    trainer: Optional["pl.Trainer"] = None
    dataloader: Optional[DataLoader] = None
    metrics: Optional[Dict[str, Any]] = None


@dataclass
class DataEvent(EventData):
    """Event data for data-related events."""

    batch: Optional[Any] = None
    batch_idx: Optional[int] = None
    dataset: Optional[Any] = None
    dataloader: Optional[DataLoader] = None


@dataclass
class TrainingEvent(EventData):
    """Event data for training-related events."""

    pl_module: Optional["pl.LightningModule"] = None
    trainer: Optional["pl.Trainer"] = None
    batch: Optional[Any] = None
    batch_idx: Optional[int] = None
    outputs: Optional[Any] = None
    optimizer_idx: Optional[int] = None
    current_epoch: Optional[int] = None
    task_id: Optional[int] = None


@dataclass
class EvaluationEvent(EventData):
    """Event data for evaluation-related events."""

    metrics: Optional[Dict[str, Any]] = None
    pl_module: Optional["pl.LightningModule"] = None
    trainer: Optional["pl.Trainer"] = None
    task_id: Optional[int] = None
    dataloader_idx: Optional[int] = None
    batch_indices: Optional[List[int]] = None


# Type variable for event handlers
E = TypeVar("E", bound=EventData)
